#!/usr/bin/python
'''
Plots average coverage of good/bad pairs after running pull_correct_edges.py
'''
import sys
import os
import h5py
import numpy as np

# parameters.
query_file = sys.argv[1]
pileup_file = sys.argv[2]

# Numpy datatypes.
names_dt = np.dtype([\
		('name', np.str_, 200),\
		('start', np.int),\
		('stop', np.int),\
		])
		
profile_dt = np.int

# open hdf5 files.
print "indexing h5 data."
h5_in = h5py.File(pileup_file, 'r')
names = h5_in['names'][0::]
profiles = h5_in['profiles'][0::]

# build id to index dictionary.
id_to_idx = {}
for i in range(len(names)):
	id_to_idx[names[i][0]] = i

# read in query file.
print "Reading queries."
fin = open(query_file, "rb")
lines = fin.readlines()
fin.close()

# get avg coverage.
total_reads = 0
avg_profile = np.zeros(50, dtype=np.int)
for line in lines:
	# tokenize.
	tmp = line.strip().split("\t")
	id1 = tmp[0]
	id2 = tmp[2]
	pos1 = int(tmp[1])
	pos2 = int(tmp[3])
	
	# look up in h5.
	ofs = names[id_to_idx[id1]][1]
	j = 0
	for i in range(ofs+pos1, ofs+pos1+50):
		avg_profile[j] += profiles[i]
		j += 1
		
	ofs = names[id_to_idx[id2]][1]
	j = 0
	for i in range(ofs+pos2, ofs+pos2+50):
		avg_profile[j] += profiles[i]
		j += 1
		
	# increment count.
	total_reads += 2
		
# close h5py file.
h5_in.close()

# Print average coverage.
for i in range(50):
	print float(avg_profile[i]) / float(total_reads)
